import numpy as np
import matplotlib.pyplot as plt
from time import clock
from search.all import *
from search.utils import infinity
from problem_agent import *
import problems
import heuristics
reload(heuristics)
reload(problems)

class beam_solution(ProblemAgent):
    
    def __init__(self, h, beam_width=infinity, max_depth=infinity):
        self.beam_width = beam_width
        self.max_depth = max_depth
        self.h = h
    
    def __repr__(self):
        return 'BS: ' + str(self.beam_width) + " " + repr(self.h)
    
    def solve(self, problem_state, time_limit = NO_LIMIT):
        '''
        This is the method called by the runner of this agent.
        It includes the code that solves the problem.
        
        @param problem_state: Initial problem state.
        @param time_limit: The time limit for this agent.
        @return: A list of ProblemActions that solves the given problem.
        '''
        search = BeamSearch(self.beam_width,self.max_depth)
        start = clock()    
        s = search.find(problem_state,self.h, time_limit)
        end = clock()
        if not s:
            return (end-start, -1)
        return (end-start, len(s))
